/*
-- IOTA Crypto Core
--
-- 2018 by Thomas Pototschnig <microengineer18@gmail.com>
-- discord: pmaxuw#8292
-- https://gitlab.com/iccfpga-rv
--
-- Permission is hereby granted, free of charge, to any person obtaining
-- a copy of this software and associated documentation files (the
-- "Software"), to deal in the Software without restriction, including
-- without limitation the rights to use, copy, modify, merge, publish,
-- distribute, sublicense, and/or sell copies of the Software, and to
-- permit persons to whom the Software is furnished to do so, subject to
-- the following conditions:
-- 
-- The above copyright notice and this permission notice shall be
-- included in all copies or substantial portions of the Software.
-- 
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
-- EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
-- MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
-- NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
-- LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
-- OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
-- WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWAR

Some code copied from Microchip's Crypto Library (See other directory)
License:

(c) 2015-2018 Microchip Technology Inc. and its subsidiaries.

Subject to your compliance with these terms, you may use Microchip software
and any derivatives exclusively with Microchip products. It is your
responsibility to comply with third party license terms applicable to your
use of third party software (including open source software) that may
accompany Microchip software.

THIS SOFTWARE IS SUPPLIED BY MICROCHIP "AS IS". NO WARRANTIES, WHETHER
EXPRESS, IMPLIED OR STATUTORY, APPLY TO THIS SOFTWARE, INCLUDING ANY IMPLIED
WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY, AND FITNESS FOR A
PARTICULAR PURPOSE. IN NO EVENT WILL MICROCHIP BE LIABLE FOR ANY INDIRECT,
SPECIAL, PUNITIVE, INCIDENTAL OR CONSEQUENTIAL LOSS, DAMAGE, COST OR EXPENSE
OF ANY KIND WHATSOEVER RELATED TO THE SOFTWARE, HOWEVER CAUSED, EVEN IF
MICROCHIP HAS BEEN ADVISED OF THE POSSIBILITY OR THE DAMAGES ARE FORESEEABLE.
TO THE FULLEST EXTENT ALLOWED BY LAW, MICROCHIP'S TOTAL LIABILITY ON ALL
CLAIMS IN ANY WAY RELATED TO THIS SOFTWARE WILL NOT EXCEED THE AMOUNT OF
FEES, IF ANY, THAT YOU HAVE PAID DIRECTLY TO MICROCHIP FOR THIS SOFTWARE.
*/

//#include "diag/Trace.h"
#include "debugprintf.h"
#include "xstatus.h"
#include "atca_execution.h"
#include "basic/atca_basic.h"
#include "hal/atca_hal.h"
#include "secure/secure.h"
#include "gpio/gpio.h"
#include "iota/conversion.h"
#include "keccak/sha3.h"


// sections MMTEXT, MMDATA, MMRODATA are protected via PMP
extern uint32_t systick;

extern ATCAIfaceCfg discoverCfg;

// needed for initializing i2c interface
extern "C" int i2c_init_axi_gpio();

MMDATA SecureSeed 		seed0(8);
MMDATA SecureSeed 		seed1(9);
MMDATA SecureSeed 		seed2(10);
MMDATA SecureSeed 		seed3(11);
MMDATA SecureSeed 		seed4(12);
MMDATA SecureSeed 		seed5(13);
MMDATA SecureSeed 		seed6(14);
MMDATA SecureSeedBIP39 	seedBip39(15);

MMDATA SecureSeed* seeds[8] = { &seed0, &seed1, &seed2, &seed3, &seed4, &seed5, &seed6, &seedBip39 };

//bool validateTrytes(const char* trytes, int len, bool checkzero);

typedef struct {
	uint8_t g_api_key[48];
	uint8_t g_slot4_key[32];
} __attribute__((packed)) keys_t;

// saved in an extra memory which is protected by PMP
static const keys_t* keys = (keys_t*) 0x84000000;

// atecc608a default config change to be atecc508a compatible
// slots 9 to 15 behaves like slot 8 ... so 8 seeds can be stored
const MMRODATA uint8_t ecc608_configdata[ATCA_ECC_CONFIG_SIZE] = { 0x01, 0x23, 0x00, 0x00, 0x00, 0x00, 0x60, 0x00, 0x04, 0x05, 0x06, 0x07, 0xEE, 0x01, 0x01,
		0x00,  //15
		0xC0, 0x00, 0x00, 0x00, 0xAF, 0x2F, 0xC4, 0x44, 0x87, 0x20, 0xC4, 0xF4, 0x8F, 0x0F, 0x0F, 0x0F,  //31, 5
		0x9F, 0x8F, 0x83, 0x64, 0xC4, 0x44, 0xC4, 0x44, 0xc4, 0x44, 0xc4, 0x44, 0xc4, 0x44, 0xc4, 0x44,  //47
		0xc4, 0x44, 0xc4, 0x44, 0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF,  //63
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x76, 0x00, 0x00, 0x00, 0x00, 0x00,  //79
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0x00, 0x40, 0x00, 0x00, 0x00, 0x00,  //95
		0x33, 0x00, 0x1C, 0x00, 0x13, 0x00, 0x1C, 0x00, 0x3C, 0x00, 0x3E, 0x00, 0x1C, 0x00, 0x33, 0x00,  //111
		0x1C, 0x00, 0x1C, 0x00, 0x1c, 0x00, 0x1c, 0x00, 0x1C, 0x00, 0x1C, 0x00, 0x1c, 0x00, 0x1c, 0x00   //127
		};
#if 0
void testdump() {
	for (int i=0;i<48;i++) {
		debugPrintf("%02x ", keys->g_api_key[i]);
	}
	debugPrintf("\n");

	for (int i=0;i<32;i++) {
		debugPrintf("%02x ", keys->g_slot4_key[i]);
	}
	debugPrintf("\n");
}
#endif


uint32_t MMTEXT SecureSeed::crc32(uint8_t* data, size_t length, uint32_t previousCrc32)
{
  uint32_t crc = ~previousCrc32;
  for (size_t i=0;i<length;i++) {
	  crc ^= data[i];
	  for (int j = 0; j < 8; j++) {
		  if (crc & 1) {
			  crc = (crc >> 1) ^ 0xEDB88320;
		  } else {
			  crc =  crc >> 1;
		  }
	  }
  }
  return ~crc;
}

bool MMTEXT SecureSeed::write(uint8_t seedBytes[48]) {
	ATCA_STATUS status;

	if (!Secure::isConfigLocked(LOCK_ZONE_DATA))
		return false;

	uint8_t write_data[64] = { 0 };

	memcpy(write_data, seedBytes, 48);

	// calculate crc32
	uint32_t crc = crc32(seedBytes, 48, 0);
	memcpy(&write_data[64-sizeof(uint32_t)], &crc, sizeof(uint32_t));

//	write_data[63] = 0x55;		// mark slot as used

	for (int block = 0; block < 2; block++) {
		status = atcab_write_enc(m_slot, block, &write_data[block*32], keys->g_slot4_key, 4);
		if (status != ATCA_SUCCESS)
			return false;
	}

	// invalidate cache
	clear();

	return true;
}

void MMTEXT SecureSeed::clear() {
	memset(m_cacheChars, 0, sizeof(m_cacheChars));
	memset(m_cacheBytes, 0, sizeof(m_cacheBytes));
	m_validSeed = false;
}
bool MMTEXT SecureSeedBIP39::loadBIP39(uint32_t bip39Path[4]) {
	// check if new bip39 path
	for (int i=0;i<4;i++) {
		if (bip39Path[i] != m_path[i]) {
			SecureSeed::setValid(false);
			break;
		}
	}
	if (SecureSeed::isCached() && SecureSeed::isValid()) {
		return true;
	}

	if (!SecureSeed::load()) {
		return false;
	}

	// if SecureSeed loaded bytes from SE, always do BIP39 calculations!
	for (int i=0;i<4;i++) {
		if (!(bip39Path[i] & 0x80000000)) {
			return false;
		}
		m_path[i] = bip39Path[i];
	}


	hdnode_from_seed((const uint8_t*) SecureSeed::getSeedBytes(), 32, "secp256k1", &m_node);
	for (int i=0;i<4;i++) {
 		if (!hdnode_private_ckd(&m_node, m_path[i])) {
			return false;
		}
	}

	uint8_t* privKey = m_node.private_key;
	uint8_t* chainCode = m_node.chain_code;

	m_keccak.init();
	m_keccak.update(&privKey[0], 32);
	m_keccak.update(&chainCode[0], 16);
	m_keccak.update(&privKey[16], 16);
	m_keccak.update(&chainCode[0], 32);
	m_keccak.final(SecureSeed::getSeedBytes());

	Conversion conv;
	conv.bytes_to_chars(SecureSeed::getSeedBytes(), SecureSeed::getSeedChars(), 48);

	SecureSeed::setValid(true);
	return true;
}
bool MMTEXT SecureSeed::load() {
	ATCA_STATUS status;

	if (m_cacheSeed && m_validSeed) {
		return true;
	}
	if (!Secure::isConfigLocked(LOCK_ZONE_DATA))
		return false;

	uint8_t read_data[64] = { 0 };

	for (int block = 0; block < 2; block++) {
		status = atcab_read_enc(m_slot, block, &read_data[block*32], keys->g_slot4_key, 4);
		if (status != ATCA_SUCCESS)
			return false;
	}

	uint32_t crc;
	memcpy(&crc, &read_data[64-sizeof(uint32_t)], sizeof(uint32_t));
	uint32_t cmp = crc32(read_data, 48, 0);

	// CRC32 wrong -> slot not valid
	if (cmp != crc) {
		return false;
	}

	memcpy(m_cacheBytes, read_data, 48);

	Conversion conv;
	conv.bytes_to_chars(m_cacheBytes, m_cacheChars, 48);

	m_validSeed = true;
	return true;
}

// only delete if caching is not allowed
void MMTEXT SecureSeed::release() {
	if (!m_cacheSeed) {
		clear();
	}
}

bool MMTEXT SecureSeed::writeRandomSeed() {
	uint8_t seedBytes[48] = {0};

	if (!Secure::generateRandomSeedBytes(seedBytes))
		return false;

	if (!write(seedBytes))
		return false;

	return true;
}



namespace Secure {
	namespace {
		bool initialized;
		bool MMTEXT lockConfigZone() {
			ATCA_STATUS status = ATCA_SUCCESS;

			if (isConfigLocked(LOCK_ZONE_CONFIG))
				return false;

			status = atcab_lock_config_zone();
			if (status != ATCA_SUCCESS)
				return false;

			if (!isConfigLocked(LOCK_ZONE_CONFIG))
				return false;

			return true;
		}

		bool MMTEXT lockDataZone() {
			ATCA_STATUS status = ATCA_SUCCESS;

			if (isConfigLocked(LOCK_ZONE_DATA))
				return false;

			status = atcab_lock_data_zone();
			if (status != ATCA_SUCCESS)
				return false;

			if (!isConfigLocked(LOCK_ZONE_DATA))
				return false;

			return true;
		}

		bool MMTEXT writeZoneConfig() {
			ATCA_STATUS status = ATCA_SUCCESS;

			if (isConfigLocked(LOCK_ZONE_CONFIG))
				return false;

			status = atcab_write_config_zone(ecc608_configdata);

			if (status != ATCA_SUCCESS)
				return false;
			return true;
		}

		bool MMTEXT writeSlot4Key(uint8_t slot4_cmp_key[32]) {
			for (int i=0;i<32;i++) {
				if (slot4_cmp_key[i] != keys->g_slot4_key[i])
					return false;
			}
			ATCA_STATUS status = ATCA_SUCCESS;

			if (!isConfigLocked(LOCK_ZONE_CONFIG))
				return false;

			status = atcab_write_zone(ATCA_ZONE_DATA, 4, 0, 0, keys->g_slot4_key, 32);
			if (status != ATCA_SUCCESS)
				return false;
			return true;
		}


		bool MMTEXT invalidateAllSlots() {
			ATCA_STATUS status = ATCA_SUCCESS;

			if (!isConfigLocked(LOCK_ZONE_DATA))
				return false;

			uint8_t write_data[64] = { 0 };

			// empty all slots
			for (int slot = 8; slot < 16; slot++) {
				for (int block = 0; block < 2; block++) {
					status = atcab_write_enc(slot, block, &write_data[block << 5], keys->g_slot4_key, 4);
					if (status != ATCA_SUCCESS)
						return false;
				}
			}

			return true;
		}
		bool MMTEXT readZone() {
			ATCA_STATUS status = ATCA_GEN_FAIL;
			uint8_t data[128];
			uint8_t serial_prefix[] = { 0x01, 0x23 };
			uint8_t slot, block, offset;
			bool locked = false;

			slot = 0;
			block = 0;
			offset = 0;

			// initialize it with recognizable data
			memset(data, 0x77, sizeof(data));

			block = 0;
			// read config zone tests
			status = atcab_read_zone(ATCA_ZONE_CONFIG, slot, block, offset, &data[block * 32], 32);
			if (status != ATCA_SUCCESS)
				return false;

			if (data[0] != serial_prefix[0] || data[1] != serial_prefix[1])
				return false;

			// read data zone tests
			// data zone cannot be read unless the data zone is locked
			status = atcab_is_locked(LOCK_ZONE_DATA, &locked);
			if (status != ATCA_SUCCESS)
				return false;

			status = atcab_read_zone(LOCK_ZONE_DATA, slot, block, offset, data, sizeof(data));
			if ((locked ? ATCA_SUCCESS : ATCA_EXECUTION_ERROR) != status)
				return false;
			return true;
		}
	}

	bool MMTEXT generateRandomSeedBytes(uint8_t seedBytes[48]) {
		ATCA_STATUS status = ATCA_SUCCESS;

		if (!Secure::isConfigLocked())
			return false;

		uint8_t rand_out[64];

		for (int i = 0; i < 2; i++) {
			status = atcab_random(&rand_out[i << 5]);
			if (status != ATCA_SUCCESS)
				return false;

		}

		// one-way-function is better
		// not using the FPGA
		Keccak384 keccak384;
		keccak384.init();
		keccak384.update(rand_out, sizeof(rand_out));
		keccak384.update(keys->g_slot4_key, 32);
		keccak384.update((uint8_t*) &systick, sizeof(uint32_t));	// some lame entropy
		keccak384.final(seedBytes);

		return true;
	}

	bool MMTEXT isConfigLocked(uint8_t zone) {
		ATCA_STATUS status = ATCA_GEN_FAIL;
		ATCAPacket packet;
		ATCACommand ca_cmd = _gDevice->mCommands;

		// build an read command
		packet.param1 = 0x00;
		packet.param2 = 0x15;
		status = atRead(ca_cmd, &packet);
		if (status != ATCA_SUCCESS) {
			return false;
		}

		status = atca_execute_command(&packet, _gDevice);
		if (status != ATCA_SUCCESS) {
			return false;
		}

		switch (zone) {
		case LOCK_ZONE_DATA:
			return packet.data[ATCA_RSP_DATA_IDX + 2] == 0;
			break;
		case LOCK_ZONE_CONFIG:
			return packet.data[ATCA_RSP_DATA_IDX + 3] == 0;
			break;
		default:
			debugPrintf("Invalid lock zone");
			break;
		}

		return false;
	}


	const uint8_t* MMTEXT getApiKey() {
		return keys->g_api_key;
	}

	bool MMTEXT initSecureElement(uint8_t slot4_cmp_key[32]) {
		// do nothing if the key isn't right
		for (int i=0;i<32;i++) {
			if (slot4_cmp_key[i] != keys->g_slot4_key[i])
				return false;
		}

		bool ret = false;
		do {
			if (!writeZoneConfig())
				break;

			if (!lockConfigZone())
				break;

			if (!writeSlot4Key(slot4_cmp_key)) {
				break;
			}

			// lock data zone
			if (!lockDataZone())
				break;

			if (!invalidateAllSlots())
				break;

			ret = true;
		} while (0);
		return ret;
	}

	bool isInitialized() {
		return initialized;
	}

	bool MMTEXT init() {
		uint8_t rev[4];
		bool ret = false;
		do {
			// a double init should be benign
			if (atcab_init(&discoverCfg) != ATCA_SUCCESS)
				break;

			// Make sure communication works initially
			if (atcab_info(rev) != ATCA_SUCCESS)
				break;

			// check if config is locked
			if (!isConfigLocked(LOCK_ZONE_CONFIG)) {
				initialized = false;
			}

			ret = true;
			initialized = true;
		} while (0);
		return ret;
	}

}


